/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stdlib.h>
#include <string.h>
#include <iostream>
#include <yhamg.h>

using namespace YHAMG;

static void SetInitialGuess(const ParMultiVec& X)
{
	int m = X.num_vectors;
	for (int j = 0; j < m; ++j)
		X(j).FillRandom();
}

int main(int argc, char* argv[])
{
	MPI_Init(&argc, &argv);

	int solver_type = 0;
	int pre_type = 0;
	int matrix_type = 0;

	int nx = 50;
	int ny = 50;
	int nz = 50;

	int bx = 1;
	int by = 1;
	int bz = 1;
	
	int Px = 1;
	int Py = 1;
	int Pz = 1;

	double cx = 1.0;
	double cy = 1.0;
	double cz = 1.0;

	int m = 4;

	int arg_index = 1;

	while (arg_index < argc)
	{
		if (strcmp(argv[arg_index], "-laplace") == 0)
			matrix_type = 0;
		else if (strcmp(argv[arg_index], "-27pt") == 0)
			matrix_type = 1;
		else if (strcmp(argv[arg_index], "-n") == 0)
		{
			nx = atoi(argv[++arg_index]);
			ny = atoi(argv[++arg_index]);
			nz = atoi(argv[++arg_index]);
		}
		if (strcmp(argv[arg_index], "-b") == 0)
		{
			bx = atoi(argv[++arg_index]);
			by = atoi(argv[++arg_index]);
			bz = atoi(argv[++arg_index]);
		}
		else if (strcmp(argv[arg_index], "-P") == 0)
		{
			Px = atoi(argv[++arg_index]);
			Py = atoi(argv[++arg_index]);
			Pz = atoi(argv[++arg_index]);
		}
		else if (strcmp(argv[arg_index], "-c") == 0)
		{
			cx = atof(argv[++arg_index]);
			cy = atof(argv[++arg_index]);
			cz = atof(argv[++arg_index]);
		}
		else if (strcmp(argv[arg_index], "-pre") == 0)
			pre_type = atoi(argv[++arg_index]);
		else if (strcmp(argv[arg_index], "-solver") == 0)
			solver_type = atoi(argv[++arg_index]);
		else if (strcmp(argv[arg_index], "-m") == 0)
			m = atof(argv[++arg_index]);
		++arg_index;
	}

	int    max_levels = 20;
	int    coarse_size = 8;
	double strength_threshold = 0.25;
	int    agg_num_levels = 1;
	int    coarsen_type = 0;
	int    interp_min_elements = 4;
	int    interp_max_elements = 4;
	double truncation_factor = 0.1;
	double sparsification_threshold = 0.01;
	int    cycle_type = 0;
	int    pre_sweeps = 1;
	int    post_sweeps = 1;
	int    coarse_sweeps = 2;
	int    relaxation_type = 3;
	double relaxation_factor = 1.0;

	int    sor_type = 2;
	double sor_factor = 1.0;

	int    gmres_restart = 10;
	int    max_iters = 500;
	double tolerance = 1.0e-08;

	bool   print_stats = 1;

	MPI_Comm comm = MPI_COMM_WORLD;

	int MPI_RANK, MPI_SIZE;
	MPI_Comm_rank(comm, &MPI_RANK);
	MPI_Comm_size(comm, &MPI_SIZE);

	if (Px * Py * Pz != MPI_SIZE) return -1;

	srand(MPI_RANK);

	if (MPI_RANK == 0)
	{
		std::cout << "\nTest Problem:\n";
		std::cout << "----------------------------------------------------------------\n";
	}

	ParOpBase* A;
	ParOpBase* P;

	int iter;
	Vec lambda, res;
	lambda.Resize(m);
	res.Resize(m);

	double setup_time;
	double solve_time;

	if (bx * by * bz == 1)
	{
		A = new ParCSR;
		
		if (matrix_type == 0)
		{
			if (MPI_RANK == 0)
			{
				std::cout << "3D Laplace Problem on a Cube\n";
				std::cout << "Global Problem Size: (" << nx * Px << ", " << ny * Py << ", " << nz * Pz << ")\n";
				std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
				std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
			}

			ParCSRGeneratorLaplace(comm, nx, ny, nz, Px, Py, Pz, cx, cy, cz)(*(ParCSR*)A);
		}
		else if (matrix_type == 1)
		{
			if (MPI_RANK == 0)
			{
				std::cout << "3D Laplace Problem with 27-point Stencil\n";
				std::cout << "Global Problem Size: (" << nx * Px << ", " << ny * Py << ", " << nz * Pz << ")\n";
				std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
				std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
			}

			ParCSRGenerator27pt(comm, nx, ny, nz, Px, Py, Pz)(*(ParCSR*)A);
		}

		((ParCSR*)A)->SetupHalo();

		if (pre_type == 0)
		{
			P = new ParCSRPrecondAMG();

			((ParCSRPrecondAMG*)P)->MaxLevels = max_levels;
			((ParCSRPrecondAMG*)P)->CoarseSize = coarse_size;
			((ParCSRPrecondAMG*)P)->StrengthThreshold = strength_threshold;
			((ParCSRPrecondAMG*)P)->AggNumLevels = agg_num_levels;
			((ParCSRPrecondAMG*)P)->CoarsenType = coarsen_type;
			((ParCSRPrecondAMG*)P)->InterpMinElements = interp_min_elements;
			((ParCSRPrecondAMG*)P)->InterpMaxElements = interp_max_elements;
			((ParCSRPrecondAMG*)P)->TruncationFactor = truncation_factor;
			((ParCSRPrecondAMG*)P)->SparsificationThreshold = sparsification_threshold;
			((ParCSRPrecondAMG*)P)->CycleType = cycle_type;
			((ParCSRPrecondAMG*)P)->PreSweeps = pre_sweeps;
			((ParCSRPrecondAMG*)P)->PostSweeps = post_sweeps;
			((ParCSRPrecondAMG*)P)->CoarseSweeps = coarse_sweeps;
			((ParCSRPrecondAMG*)P)->RelaxationType = relaxation_type;
			((ParCSRPrecondAMG*)P)->RelaxationFactor = relaxation_factor;
			((ParCSRPrecondAMG*)P)->PrintStats = print_stats;
		}
		else if (pre_type == 1)
		{
			P = new ParCSRPrecondILUAMG;

			((ParCSRPrecondILUAMG*)P)->MaxLevels = max_levels;
			((ParCSRPrecondILUAMG*)P)->CoarseSize = coarse_size;
			((ParCSRPrecondILUAMG*)P)->StrengthThreshold = strength_threshold;
			((ParCSRPrecondILUAMG*)P)->AggNumLevels = agg_num_levels;
			((ParCSRPrecondILUAMG*)P)->CoarsenType = coarsen_type;
			((ParCSRPrecondILUAMG*)P)->InterpMinElements = interp_min_elements;
			((ParCSRPrecondILUAMG*)P)->InterpMaxElements = interp_max_elements;
			((ParCSRPrecondILUAMG*)P)->TruncationFactor = truncation_factor;
			((ParCSRPrecondILUAMG*)P)->SparsificationThreshold = sparsification_threshold;
			((ParCSRPrecondILUAMG*)P)->CycleType = cycle_type;
			((ParCSRPrecondILUAMG*)P)->PreSweeps = pre_sweeps;
			((ParCSRPrecondILUAMG*)P)->PostSweeps = post_sweeps;
			((ParCSRPrecondILUAMG*)P)->CoarseSweeps = coarse_sweeps;
			((ParCSRPrecondILUAMG*)P)->PrintStats = print_stats;
		}
		else if (pre_type == 2)
			P = new ParCSRPrecondJacobi;
		else if (pre_type == 3)
			P = new ParCSRPrecondSOR(sor_type, sor_factor);
		else if (pre_type == 4)
			P = new ParCSRPrecondILU;

		MPI_Barrier(comm);

		if (MPI_RANK == 0)
		{
			setup_time = -MPI_Wtime();

			if (pre_type == 0)
				std::cout << "\nAMG Setup Phase:\n";
			else if (pre_type == 1)
				std::cout << "\nILUAMG Setup Phase:\n";
			else if (pre_type == 2)
				std::cout << "\nJacobi Setup Phase:\n";
			else if (pre_type == 3)
				std::cout << "\nSOR Setup Phase:\n";
			else if (pre_type == 4)
				std::cout << "\nILU Setup Phase:\n";
			
			std::cout << "----------------------------------------------------------------\n";
		}

		((ParCSRPrecondBase*)P)->Setup(*(ParCSR*)A);

		MPI_Barrier(comm);
		
		if (MPI_RANK == 0)
		{
			setup_time += MPI_Wtime();

			std::cout << "Setup Time: " << setup_time << "\n";
		}
	}
	else
	{
		A = new ParBSR;
		
		if (matrix_type == 0)
		{
			if (MPI_RANK == 0)
			{
				std::cout << "3D Laplace Problem on a Cube\n";
				std::cout << "Global Problem Size: (" << nx * Px * bx << ", " << ny * Py * by << ", " << nz * Pz * bz << ")\n";
				std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
				std::cout << "Nonzero Block Size: (" << bx << ", " << by << ", " << bz << ")\n";
				std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
			}

			ParBSRGeneratorLaplace(comm, nx, ny, nz, bx, by, bz, Px, Py, Pz, cx, cy, cz)(*(ParBSR*)A);
		}
		else if (matrix_type == 1)
		{
			if (MPI_RANK == 0)
			{
				std::cout << "3D Laplace Problem with 27-point Stencil\n";
				std::cout << "Global Problem Size: (" << nx * Px * bx << ", " << ny * Py * by << ", " << nz * Pz * bz << ")\n";
				std::cout << "Local Domain Size: (" << nx << ", " << ny << ", " << nz << ")\n";
				std::cout << "Nonzero Block Size: (" << bx << ", " << by << ", " << bz << ")\n";
				std::cout << "MPI Process Size: (" << Px << ", " << Py << ", " << Pz << ")\n";
			}

			ParBSRGenerator27pt(comm, nx, ny, nz, bx, by, bz, Px, Py, Pz)(*(ParBSR*)A);
		}

		((ParBSR*)A)->SetupHalo();

		if (pre_type == 0)
		{
			P = new ParBSRPrecondAMG();

			((ParBSRPrecondAMG*)P)->MaxLevels = max_levels;
			((ParBSRPrecondAMG*)P)->CoarseSize = coarse_size;
			((ParBSRPrecondAMG*)P)->StrengthThreshold = strength_threshold;
			((ParBSRPrecondAMG*)P)->AggNumLevels = agg_num_levels;
			((ParBSRPrecondAMG*)P)->CoarsenType = coarsen_type;
			((ParBSRPrecondAMG*)P)->InterpMinElements = interp_min_elements;
			((ParBSRPrecondAMG*)P)->InterpMaxElements = interp_max_elements;
			((ParBSRPrecondAMG*)P)->TruncationFactor = truncation_factor;
			((ParBSRPrecondAMG*)P)->SparsificationThreshold = sparsification_threshold;
			((ParBSRPrecondAMG*)P)->CycleType = cycle_type;
			((ParBSRPrecondAMG*)P)->PreSweeps = pre_sweeps;
			((ParBSRPrecondAMG*)P)->PostSweeps = post_sweeps;
			((ParBSRPrecondAMG*)P)->CoarseSweeps = coarse_sweeps;
			((ParBSRPrecondAMG*)P)->RelaxationType = relaxation_type;
			((ParBSRPrecondAMG*)P)->RelaxationFactor = relaxation_factor;
			((ParBSRPrecondAMG*)P)->PrintStats = print_stats;
		}
		else if (pre_type == 1)
		{
			P = new ParBSRPrecondILUAMG;

			((ParBSRPrecondILUAMG*)P)->MaxLevels = max_levels;
			((ParBSRPrecondILUAMG*)P)->CoarseSize = coarse_size;
			((ParBSRPrecondILUAMG*)P)->StrengthThreshold = strength_threshold;
			((ParBSRPrecondILUAMG*)P)->AggNumLevels = agg_num_levels;
			((ParBSRPrecondILUAMG*)P)->CoarsenType = coarsen_type;
			((ParBSRPrecondILUAMG*)P)->InterpMinElements = interp_min_elements;
			((ParBSRPrecondILUAMG*)P)->InterpMaxElements = interp_max_elements;
			((ParBSRPrecondILUAMG*)P)->TruncationFactor = truncation_factor;
			((ParBSRPrecondILUAMG*)P)->SparsificationThreshold = sparsification_threshold;
			((ParBSRPrecondILUAMG*)P)->CycleType = cycle_type;
			((ParBSRPrecondILUAMG*)P)->PreSweeps = pre_sweeps;
			((ParBSRPrecondILUAMG*)P)->PostSweeps = post_sweeps;
			((ParBSRPrecondILUAMG*)P)->CoarseSweeps = coarse_sweeps;
			((ParBSRPrecondILUAMG*)P)->PrintStats = print_stats;
		}
		else if (pre_type == 2)
			P = new ParBSRPrecondJacobi;
		else if (pre_type == 3)
			P = new ParBSRPrecondSOR(sor_type, sor_factor);
		else if (pre_type == 4)
			P = new ParBSRPrecondILU;

		MPI_Barrier(comm);

		if (MPI_RANK == 0)
		{
			setup_time = -MPI_Wtime();

			if (pre_type == 0)
				std::cout << "\nAMG Setup Phase:\n";
			else if (pre_type == 1)
				std::cout << "\nILUAMG Setup Phase:\n";
			else if (pre_type == 2)
				std::cout << "\nJacobi Setup Phase:\n";
			else if (pre_type == 3)
				std::cout << "\nSOR Setup Phase:\n";
			else if (pre_type == 4)
				std::cout << "\nILU Setup Phase:\n";
			
			std::cout << "----------------------------------------------------------------\n";
		}

		((ParBSRPrecondBase*)P)->Setup(*(ParBSR*)A);

		MPI_Barrier(comm);
		
		if (MPI_RANK == 0)
		{
			setup_time += MPI_Wtime();

			std::cout << "Setup Time: " << setup_time << "\n";
		}
	}

	ParMultiVec X(comm);

	X.Allocate(m, A->OutSize());

	SetInitialGuess(X);

	MPI_Barrier(comm);

	if (MPI_RANK == 0)
	{
		solve_time = -MPI_Wtime();

		if (solver_type == 0)
			std::cout << "\nLOBPCG Solve Phase:\n";
		std::cout << "----------------------------------------------------------------\n";
	}

	if (solver_type == 0)
		ParEigenSolverLOBPCG(ParOpComp(-1, *A, 0.0), *P, max_iters, tolerance, print_stats)(X, lambda, iter, res);

	MPI_Barrier(comm);

	if (MPI_RANK == 0)
	{
		solve_time += MPI_Wtime();

		std::cout << "Solve Time: " << solve_time << "\n";
		std::cout << "Time/Iteration: " << solve_time / iter << "\n";
	}

	if (MPI_RANK == 0)
	{
		std::cout << "\n----------------------------------------------------------------\n";
		std::cout << "Total Time: " << setup_time + solve_time << "\n";
	}

	if (MPI_RANK == 0)
	{
		std::cout << "\n";
		for (int j = 0; j < m; ++j)
			std::cout <<"LAMBDA #" << j + 1 << " = "<< lambda.values[j] << "\n";
	}

	delete A;
	delete P;

	MPI_Finalize();

	return 0;
}
